import paramiko
import getpass
import time

# SSH 连接配置
HOST = '10.21.162.157'  # 替换为你的服务器IP
PORT = 22
USERNAME = 'root'  # 替换为你的用户名

# 获取密码
PASSWORD = getpass.getpass(f"Enter SSH password for {USERNAME}@{HOST}: ")

# 创建SSH客户端
client = paramiko.SSHClient()
client.set_missing_host_key_policy(paramiko.AutoAddPolicy())

try:
    # 建立SSH连接
    print(f"Connecting to {HOST}...")
    client.connect(HOST, PORT, USERNAME, PASSWORD)
    
    # 创建交互式shell
    shell = client.invoke_shell()
    
    # 设置非阻塞模式
    shell.setblocking(False)
    
    # 等待shell准备
    time.sleep(1)
    
    # 读取初始欢迎信息
    if shell.recv_ready():
        print(shell.recv(1024).decode(), end='')
    
    # 主循环
    while True:
        command = input(" ")  # 等待用户输入命令
        
        if command.lower() == 'exit':
            break
            
        # 发送命令
        shell.send(command + '\n')
        
        # 等待命令执行
        time.sleep(0.1)
        
        # 读取输出
        output = ""
        while True:
            if shell.recv_ready():
                data = shell.recv(4096).decode()
                output += data
                # 实时打印输出
                print(data, end='')
                
                # 检查是否需要用户输入（如sudo密码提示）
                if "[sudo]" in data or "password" in data.lower():
                    password = getpass.getpass("Enter sudo password: ")
                    shell.send(password + '\n')
                    time.sleep(0.5)  # 等待密码处理
            else:
                # 如果没有更多数据且检测到命令提示符，结束读取
                if output.endswith('$ ') or output.endswith('# '):
                    break
                time.sleep(0.1)
    
    print("Session ended.")
    
finally:
    # 关闭连接
    client.close()
    print("SSH connection closed.")